import torch
import torch.nn as nn
import math
import torch.nn.functional as F


class BertEmbedding(nn.Module):
    def __init__(self, config):
        super(BertEmbedding, self).__init__()
        self.word_embedding = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_idx)
        self.position_embedding = nn.Embedding(config.max_position, config.hidden_size)
        self.token_type_embedding = nn.Embedding(config.type_size, config.hidden_size)
        self.layer_norm = nn.LayerNorm(config.hidden_size)
        self.dropout = nn.Dropout(config.hidden_dropout)
        self.device = torch.device(config.device)
        self.register_buffer('position_ids', torch.arange(config.max_position).expand(1, -1))

    def forward(self, input_ids, token_type_ids=None):
        input_shape = input_ids.size()
        input_embedding = self.word_embedding(input_ids)
        if token_type_ids is None:
            token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=self.device)
        token_type_embedding = self.token_type_embedding(token_type_ids)
        position_embedding = self.position_embedding(self.position_ids)
        embedding = input_embedding + token_type_embedding + position_embedding
        embedding = self.dropout(self.layer_norm(embedding))
        return embedding


class BertSelfAttention(nn.Module):
    def __init__(self, config):
        super(BertSelfAttention, self).__init__()
        self.attention_heads = config.attention_heads
        self.attention_head_size = config.hidden / self.attention_heads
        self.all_head_size = self.attention_head_size * self.attention_heads
        self.query = nn.Linear(config.hidden, self.all_head_size)
        self.key = nn.Linear(config.hidden, self.all_head_size)
        self.value = nn.Linear(config.hidden, self.all_head_size)
        self.dropout = nn.Dropout(config.attention_dropout)

    def transpose_for_score(self, x):
        batch, seq = x[0], x[1]
        x = x.view(batch, seq, self.attention_heads, self.attention_head_size)
        return x.permute(0, 2, 1, 3)

    def forward(self, hidden):
        key_layer = self.transpose_for_score(self.key(hidden))
        value_layer = self.transpose_for_score(self.value(hidden))
        query_layer = self.transpose_for_score(self.query(hidden))
        attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2)) / math.sqrt(self.attention_head_size)
        attention_probs = nn.Softmax(dim=-1)(attention_scores)
        attention_probs = self.dropout(attention_probs)
        context_layer = torch.matmul(attention_probs, value_layer)
        context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
        context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
        context_layer = context_layer.view(*context_layer_shape)
        return context_layer, attention_probs


class BertSelfOutput(nn.Module):
    def __init__(self, config):
        super(BertSelfOutput, self).__init__()
        self.dense = nn.Linear(config.hidden, config.hidden)
        self.dropout = nn.Dropout(config.hidden_dropout)
        self.layer_norm = nn.LayerNorm(config.hidden_dropout)

    def forward(self, hidden, input_tensor):
        hidden = self.dropout(self.dense(hidden))
        return self.layer_norm(hidden + input_tensor)


class BertAttention(nn.Module):
    def __init__(self, config):
        super(BertAttention, self).__init__()
        self.attention = BertSelfAttention(config)
        self.output = BertSelfOutput(config)

    def forward(self, hidden):
        context, attention = self.attention(hidden)
        output = self.output(context, hidden)
        return output, attention


class BertIntermediate(nn.Module):
    def __init__(self, config):
        super(BertIntermediate, self).__init__()
        self.dense = nn.Linear(self.hidden, self.intermediate)

    def forward(self, hidden):
        return F.gelu(self.dense(hidden))


class BertOutput(nn.Module):
    def __init__(self, config):
        super(BertOutput, self).__init__()
        self.dense = nn.Linear(config.intermediate, config.hidden)
        self.dropout = nn.Dropout(config.hidden_dropout)
        self.layer_norm = nn.LayerNorm(config.hidden)

    def forward(self, hidden, input_tensor):
        return self.layer_norm(self.dropout(self.dense(hidden)) + input_tensor)


class BertLayer(nn.Module):
    def __init__(self, config):
        super(BertLayer, self).__init__()
        self.attention = BertAttention(config)
        self.intermediate = BertIntermediate(config)
        self.output = BertOutput(config)

    def forward(self, hidden):
        output, attention = self.attention(hidden)
        intermediate_output = self.intermediate(output)
        output = self.output(intermediate_output, output)
        return output, attention


class BertEncoder(nn.Module):
    def __init__(self, config):
        super(BertEncoder, self).__init__()


